/*
 * Copyright 2008 Neil Bacon
 * Distributed under the terms of the Mozilla Public License version 1.1
 * or the GNU General Public License version 3.
 */

package nb.xadd.operator.impl;

import static java.lang.Double.isNaN;

import java.util.HashMap;
import java.util.Map;

import nb.cache.CacheMap;
import nb.util.Bounds;
import nb.util.Utils;
import nb.xadd.XADDRuntimeException;
import nb.xadd.node.Node;
import nb.xadd.node.NodeFactory;
import nb.xadd.node.impl.AbstractDecisionNode;
import nb.xadd.node.impl.BooleanDecisionNode;
import nb.xadd.node.impl.ConstantValueNode;
import nb.xadd.node.impl.LinearValueNode;
import nb.xadd.node.impl.RealEqualDecisionNode;
import nb.xadd.node.impl.RealGreaterOrEqualDecisionNode;
import nb.xadd.operator.BinaryOperator;
import nb.xadd.operator.OperatorFactory;
import nb.xadd.operator.UnaryOperator;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * Replace LinearValueNodes with a Piecewise Constant Function.
 * All terminal nodes must be ConstantValueNodes or LinearValueNodes or a MDPRuntimeException is thrown.
 * TODO: handle PolynomialValueNodes.
 */
public class PieceWiseConstantApproximation implements UnaryOperator {

    private Log log = LogFactory.getLog(getClass());

    private final NodeFactory nodeFactory;
    private final CacheMap<UnaryOperatorCacheKey, Node> unaryOperatorCache;
    private final Type type;
    private final double errorLimit;
    private final int pieceCountLimit;
    private final int hashCode;
    
    private final BinaryOperator times;
    
    private final Node zero;
    private final Node one;

    protected PieceWiseConstantApproximation(NodeFactory nodeFactory, OperatorFactory operatorFactory,
            CacheMap<UnaryOperatorCacheKey, Node> unaryOperatorCache, Type type, double errorLimit, int pieceCountLimit) {
        super();
        this.nodeFactory = nodeFactory;
        this.unaryOperatorCache = unaryOperatorCache;
        this.type = type;
        this.errorLimit = errorLimit;
        this.pieceCountLimit = pieceCountLimit;
        hashCode = calcHashCode();
        
        times = operatorFactory.getMultiplication();
        
        zero = nodeFactory.getConstantValueNode(0.0);
        one = nodeFactory.getConstantValueNode(1.0);
    }

    public Node execute(Node node) {
        UnaryOperatorCacheKey key = new UnaryOperatorCacheKey(this, node);
        Node result = unaryOperatorCache.get(key);
        if (result == null) {
            result = walk(node, new HashMap<Integer, Bounds>());
            unaryOperatorCache.put(key, result);
        }
        return result;
    }

    private Node walk(Node node, Map<Integer, Bounds> boundsMap) {
        Node result = null;
        if (node instanceof RealGreaterOrEqualDecisionNode) {
            RealGreaterOrEqualDecisionNode dn = (RealGreaterOrEqualDecisionNode) node;
            Bounds bounds = Bounds.getBounds(boundsMap, dn.getRealVarIndex());
            if (log.isDebugEnabled()) {
                log.debug("walk: dn x" + dn.getRealVarIndex() + " >= " + dn.getTestValue() + ", Bounds = " + bounds);
            }

            double min = bounds.min;
            bounds.min = dn.getTestValue(); // descendants get bounds modified according to this test
            Node trueChild = walk(dn.getTrueChild(), boundsMap);
            bounds.min = min;               // restore current bounds after processing descendants

            double max = bounds.max;
            bounds.max = dn.getTestValue();
            Node falseChild = walk(dn.getFalseChild(), boundsMap);
            bounds.max = max;

            if (isNaN(bounds.min)) {
                result = nodeFactory.getRealGreaterOrEqualDecisionNode(dn.getRealVarIndex(), dn.getTestValue(), trueChild, falseChild);
            } else {
                double xRange = dn.getTestValue() - bounds.min;
                double yRange = Math.abs(xRange * getCoeff(falseChild, dn.getRealVarIndex()));
                double numPieces = Math.ceil(type == Type.MID ? yRange / (errorLimit * 2.0) : yRange / errorLimit);
                if (numPieces > pieceCountLimit) {
                    if (log.isWarnEnabled()) {
                        log.warn("walk: error bounds would require " + numPieces + " pieces, but we limit to " + pieceCountLimit);
                    }
                    numPieces = pieceCountLimit;
                }
                double inc = xRange / numPieces;
                if (log.isDebugEnabled()) {
                    log.debug("walk: xRange " + xRange + ", yRange = " + yRange + ", numPieces = " + numPieces + ", inc = " + inc);
                }
                Node evalFalseChild = evaluate(falseChild, dn.getRealVarIndex(), dn.getTestValue(), inc);
                RealGreaterOrEqualDecisionNode newDn = (RealGreaterOrEqualDecisionNode)
                    nodeFactory.getRealGreaterOrEqualDecisionNode(dn.getRealVarIndex(), dn.getTestValue(), trueChild, evalFalseChild);
                if (log.isDebugEnabled()) {
                    log.debug("walk: newDn " + newDn);
                }
                result = createDecisionNode(newDn, falseChild, bounds, inc);
            }
        } else if (node instanceof RealEqualDecisionNode) {
            RealEqualDecisionNode dn = (RealEqualDecisionNode) node;
            Node trueChild = walk(dn.getTrueChild(), boundsMap);
            Node falseChild = walk(dn.getFalseChild(), boundsMap);
            result = nodeFactory.getRealEqualDecisionNode(dn.getRealVarIndex(), dn.getTestValue(), trueChild, falseChild);
        } else if (node instanceof BooleanDecisionNode) {
            BooleanDecisionNode dn = (BooleanDecisionNode) node;
            Node trueChild = walk(dn.getTrueChild(), boundsMap);
            Node falseChild = walk(dn.getFalseChild(), boundsMap);
            result = nodeFactory.getBooleanDecisionNode(dn.getBooleanVarIndex(), trueChild, falseChild);
        } else if (node instanceof LinearValueNode) {
            log.debug("walk: boundsMap = " + boundsMap);
            result = node;
        } else if (node instanceof ConstantValueNode) {
            result = node;
        } else {
            throw new XADDRuntimeException("PieceWiseConstantApproximation can't handle Node type " + node.getClass().getName());
        }
        return result;
    }

    /**
     * Get linear coefficient for realVarIndex in n.
     * @param n
     * @param realVarIndex
     * @return coefficient for realVarIndex in the first LinearValueNode we encounter in n
     * or Double.NaN if none.
     */
    private double getCoeff(Node n, int realVarIndex) {
        if (n instanceof LinearValueNode) {
            LinearValueNode ln = (LinearValueNode) n;
            return realVarIndex < ln.getCoeff().length ? ln.getCoeff()[realVarIndex] : 0.0;
        } else if (n instanceof AbstractDecisionNode) {
            AbstractDecisionNode dn = (AbstractDecisionNode) n;
            double c = getCoeff(dn.getTrueChild(), realVarIndex);
            return isNaN(c) ? getCoeff(dn.getFalseChild(), realVarIndex) : c;
        }
        return Double.NaN;
    }

    /**
     * Bind value to variable realVarIndex.
     * @param n
     * @param realVarIndex
     * @param value
     * @return
     */
    Node bind(Node n, int realVarIndex, double value) {
        Node n2 = nodeFactory.getRealEqualDecisionNode(realVarIndex, value, one, zero);
        return ((AbstractDecisionNode) n.apply(times, n2)).getTrueChild();
    }    
    
    /**
     * Evaluate falseChild with realVar[realVarIndex] bound to value depending on type.
     * @param falseChild
     * @param realVarIndex
     * @param testValue
     * @param inc
     * @return
     */
    private Node evaluate(Node falseChild, int realVarIndex, double testValue, double inc) {
        double value = type == Type.UPPER ? testValue
                : type == Type.LOWER ? testValue - inc
                        : testValue - inc / 2.0;
        Node result = bind(falseChild, realVarIndex, value);
        if (log.isDebugEnabled()) {
            log.debug("evaluate: bind x" + realVarIndex + " = " + value + ", result =\n" + result);
        }
        return result;
    }

    private RealGreaterOrEqualDecisionNode createDecisionNode(RealGreaterOrEqualDecisionNode trueChild, Node falseChild, Bounds bounds, double inc) {
        while (trueChild.getTestValue() - inc > bounds.min) {
            double testValue = trueChild.getTestValue() - inc;
            int realVarIndex = trueChild.getRealVarIndex();
            Node evalFalseChild = evaluate(falseChild, realVarIndex, testValue, inc);
            trueChild = (RealGreaterOrEqualDecisionNode) nodeFactory.getRealGreaterOrEqualDecisionNode(realVarIndex, testValue, trueChild, evalFalseChild);
            if (log.isDebugEnabled()) {
                log.debug("createDecisionNode: result " + trueChild);
            }
        }
        return trueChild;
    }

    protected int calcHashCode() {
        final int prime = 61;
        int result = 1;
        result = prime * result + type.hashCode();
        result = prime * result + Utils.hashCode(errorLimit);
        result = prime * result + pieceCountLimit;
        return result;
    }

    @Override
    public int hashCode() {
        return hashCode;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (obj == null || getClass() != obj.getClass())
            return false;
        PieceWiseConstantApproximation other = (PieceWiseConstantApproximation) obj;
        return errorLimit == other.errorLimit;
    }


    // Nested classes...

    /**
     * Type of approximation.
     */
    public static enum Type {
        /** Lower bound approximation. */
        LOWER,
        /** Best approximation. */
        MID,
        /** Upper bound approximation. */
        UPPER
    }

}
